'''
Created on Nov 21, 2012
@author: rafa
'''

import numpy as np
import nipype.interfaces.io as nio
import nipype.interfaces.fsl as fsl
import nipype.interfaces.spm as spm
import nipype.interfaces.utility as util
import nipype.pipeline.engine as pe
import nipype.algorithms.modelgen as model
import os

import nipype.pipeline.plugins

fsl.FSLCommand.set_default_output_type('NIFTI')

#Data set up
data = './Inputs'
os.chdir(data)
data_dir = os.path.abspath('.')
subject_list = ['Sub_1','Sub_2']

#info=dict(func=[['subject_id',['EO',]]],struct=[['subject_id', 'struct']])
#infosource=pe.Node(interface=util.IdentityInterface(fields=['subject_id']),name="infosource")
#infosource.iterables=('subject_id', subject_list)

datasource=pe.Node(interface=nio.DataGrabber(infields=['subject_id'], outfields=['func', 'struct']), name='datasource')
datasource.inputs.base_directory=data_dir
datasource.inputs.template='%s/%s'
datasource.inputs.field_template = dict(func=('%s/BOLD/%s_BOLD.nii'),struct=('%s/T1/o_%s_T1.nii'))
datasource.inputs.template_args = dict(func=[['subject_id', 'subject_id']],struct=[['subject_id', 'subject_id']])
datasource.iterables = ('subject_id', subject_list)

#structural processing
extract_ref = pe.Node(interface=fsl.ExtractROI(t_min=42, t_size=1), name='extractref')

#nosestrip = pe.Node(interface=fsl.BET(frac=3.0), name='nosestrip')
skullstrip = pe.Node(interface=fsl.BET(mask = True), name='stripstruct')
refskullstrip = pe.Node(interface=fsl.BET(mask = True), name = 'stripref')
coregister = pe.Node(interface=fsl.FLIRT(dof=6), name = 'coregister')


#Functional processing
motion_correct = pe.Node(interface=fsl.MCFLIRT(save_plots = True), name = 'realign')
func_skullstrip = pe.Node(interface=fsl.BET(functional = True), name='stripfunc')
smoother= pe.Node(interface=spm.Smooth(), name='smooth')
smoother.iterables=('fwhm', [0,3,6])

####
#specific to CSF method
getCSFmask = pe.Node(interface=fsl.FAST(no_pve=True, segments=True), name='segment')
applyReg2CSFmask = pe.Node(interface=fsl.ApplyXfm(apply_xfm=True), name='applyreg2csfmask')
threshCSFseg = pe.Node(interface=fsl.ImageMaths(op_string = ' -thr .90 -uthr 1 -bin '), name = 'threshcsfsegmask')

#Get CSF time series
avgCSF = pe.Node(interface = fsl.ImageMeants(), name='extractcsfts')


#function for taking first volume
def pickfirst(files):
    return files[0]

def picksecond(files):
    return files[1]

#first workflow for csf
csffilter = pe.Workflow(name='csffilter')
csffilter.connect([(extract_ref, motion_correct, [('roi_file', 'ref_file')]),
                         (extract_ref, refskullstrip, [('roi_file', 'in_file')]),
                         (skullstrip, getCSFmask, [('out_file', 'in_files')]),
                         (skullstrip, coregister, [('mask_file','in_file')]),
                         (refskullstrip, coregister, [('out_file','reference')]),
                         (motion_correct, smoother, [('out_file','in_files')]),
                         (smoother, func_skullstrip, [('smoothed_files','in_file')]),
                         (getCSFmask, applyReg2CSFmask, [(('tissue_class_files', pickfirst), 'in_file')]),
                         (refskullstrip, applyReg2CSFmask, [('out_file', 'reference')]),
                         (coregister, applyReg2CSFmask, [('out_matrix_file', 'in_matrix_file')]),
                         (applyReg2CSFmask, threshCSFseg,[('out_file','in_file')]),
                         (func_skullstrip, avgCSF,[('out_file','in_file')]),
                         (threshCSFseg, avgCSF, [('out_file','mask')]),
                         ])


# model side
modelfit = pe.Workflow(name='modelfit')
modelspec=pe.Node(interface=model.SpecifyModel(), name="modelspec")
level1design=pe.Node(interface=fsl.Level1Design(), name="fsfdesign")
modelgen=pe.Node(interface=fsl.FEATModel(), name="modelgen")
modelestimate=pe.Node(interface=fsl.FILMGLS(), name="modelestimate")

modelfit.connect([(modelspec, level1design, [('session_info','session_info')]),
                                            (level1design, modelgen,[('fsf_files', 'fsf_file'),
                                            ('ev_files', 'ev_files')]),
                                            (modelgen, modelestimate, [('design_file','design_file')]),
                ])


def subjectinfo(meantsfile):
    import numpy as np
    from  nipype.interfaces.base import Bunch
    ts = np.loadtxt(meantsfile)
    output = [Bunch(regressor_names=['MeanIntensity'], regressors=[ts.tolist()])]
    return output

hpcutoff = np.Inf
TR = 3. 

modelfit.inputs.modelspec.input_units='secs'
modelfit.inputs.modelspec.time_repetition=TR
modelfit.inputs.modelspec.high_pass_filter_cutoff=hpcutoff
modelfit.inputs.fsfdesign.interscan_interval=TR
modelfit.inputs.fsfdesign.bases={'none':None}
modelfit.inputs.fsfdesign.model_serial_correlations=False
modelfit.inputs.modelestimate.autocorr_noestimate=True
bandPassFilterData = pe.Node(interface=fsl.ImageMaths(op_string=' -bptf 128 12.5 '), name='bandpassfiltermcdata_fslmaths')

##Complete workflow
l1pipeline = pe.Workflow(name = "resting3")
l1pipeline.base_dir = './Inputs'
l1pipeline.connect([(datasource, csffilter, [('struct', 'stripstruct.in_file'), 
                                                                    ('func', 'realign.in_file'),
                                                                    ('func', 'extractref.in_file'),
                                                                    ]),
                               (csffilter, modelfit, [('stripfunc.out_file', 'modelspec.functional_runs'),
                                                      ('realign.par_file', 'modelspec.realignment_parameters'),
                                                      (('extractcsfts.out_file', subjectinfo), 'modelspec.subject_info'),
                                                      ('stripfunc.out_file', 'modelestimate.in_file')
                                                      ]),
                              (modelfit, bandPassFilterData, [('modelestimate.residual4d', 'in_file')])
                     ])


if __name__ == '__main__':
    l1pipeline.run()
    l1pipeline.write_graph()
